import os
import re
import sys

from data import setup

root = os.path.dirname(os.path.dirname(__file__))
sys.path.append(root)

def get_flops(config_path, code_dir='', task='', readme='', export=False, device='cpu'):

    if readme:
        checkpoints = [c for c in readme.strip().split('\n') if f'/{config_path}.pt' in c]
        params, flops = checkpoints[0].split('|')[-3:-1]
        flops = re.sub(' *\**', '', flops)
        params = re.sub(' *\** *', '', params)
        return flops+' GFLOPs', params+ ' M'
    else:
        torch_utils_path = os.path.join(
            code_dir, 'utils', 'torch_utils.py')
        torch_utils_content = open(torch_utils_path).read()

        torch_utils_content = re.sub(
            '(fs}"\))', r'\1\n    return fs, n_p', torch_utils_content)

        open(torch_utils_path, 'w').write(torch_utils_content)

        try:
            from downloads.yolov5.models import yolo
            from downloads.yolov5.utils import torch_utils
        except:
            return

        if task == 'detection':
            model = yolo.DetectionModel(config_path).to(device)
        elif task == 'classification':
            model = yolo.ClassificationModel(config_path).to(device)

        if export:
            return model
        else:
            flops, params = torch_utils.model_info(model)
            return flops[2:], f'{round(params / 1e6, 1)} M'

def run(task='', data_dir='', github_dir=''):

    git_dir, code_dir = setup(task=task, data_dir=data_dir, github_dir=github_dir)

    try:
        from _configs import configs
        from util_my import save_config
    except:
        configs = {}

    # TODO yamls torchvisions
    # yamls = sorted(glob.glob(os.path.join(code_dir, 'models', '**', '**', 'yolo**.yaml'), recursive=True))
    readme = open(os.path.join(git_dir, 'README.md'), encoding='utf-8').read()

    # detection
    detection = readme[readme.index('Pretrained Checkpoints'):]
    detection = detection[detection.index('| Model'):]
    detection = detection[:detection.index('<details>')]
    det_ids = re.findall('(?<=download/v\d.\d/).*?(?=.pt)', detection)

    # classification
    classification = readme[readme.index('Classification'):]
    classification = classification[classification.index('| Model'):]
    classification = classification[:classification.index('<details>')]
    cls_ids = re.findall('(?<=download/v\d.\d/).*?(?=.pt)', classification)

    # segmentation
    segmentation = readme[readme.index('Segmentation'):]
    segmentation = segmentation[segmentation.index('| Model'):]
    segmentation = segmentation[:segmentation.index('<details>')]
    seg_ids = re.findall('(?<=download/v\d.\d/).*?(?=.pt)', segmentation)

    cfg_ids = det_ids + cls_ids + seg_ids

    import torchvision

    for key, value in torchvision.models.__dict__.items():
        if 'function' in str(type(value)) and key not in cls_ids:
            cfg_ids.append(key)

    import tqdm
    for config_path in tqdm.tqdm(cfg_ids):

        cfg_id = str(config_path).replace(os.path.join(code_dir, 'models'), '')[1:] \
            if os.path.exists(config_path) else config_path
        cfg_id = cfg_id.replace('.yaml', '').replace(f'hub{os.path.sep}', '')

        readme_ = readme if cfg_id in det_ids or cfg_id in cls_ids else ''

        if cfg_id in det_ids:
            task = 'detection'
        if cfg_id in cls_ids:
            task = 'classification'
            cfg_id = cfg_id if 'yolov5' in cfg_id else f'yolov5-{cfg_id}'
        if cfg_id in seg_ids:
            task = 'segmentation'

        if cfg_id not in configs:
            configs[task][cfg_id] = {}

        width = height = 640 if cfg_id in det_ids else 224

        flops_id = f'flops_{width}x{height}'

        if flops_id not in configs[task][cfg_id] or \
                'params' not in configs[task][cfg_id]:

            print('\n', flops_id)

            try:
                flops, params = get_flops(config_path, code_dir=code_dir, task=task, readme=readme_)
            except:
                del configs[task][cfg_id]
                continue

            configs[task][cfg_id]['params'] = params.replace('<br>-', '')
            configs[task][cfg_id][flops_id] = flops.replace('<br>-', '')

            save_config(configs, task, root)

if __name__ == '__main__':
    pass